from matplotlib import pyplot as plt
import numpy as np
from GGENet_torch_train import get_length_classes
import itertools
def plot_confusion_matrix(confused_matrix,num_classes,normalize=False, title='State transition matrix', cmap=plt.cm.Blues):
    classes = [i for i in range(num_classes)]

    matrix = np.zeros((num_classes,num_classes))
    for i in range(len(confused_matrix[0])):
        x = confused_matrix[0][i]
        y = confused_matrix[1][i]
        matrix[x][y] += 1

    plt.figure()
    plt.imshow(matrix, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(num_classes)
    plt.xticks(tick_marks, classes, rotation=90)
    plt.yticks(tick_marks, classes)

    plt.axis("equal")

    ax = plt.gca()
    left, right = plt.xlim()
    ax.spines['left'].set_position(('data', left))
    ax.spines['right'].set_position(('data', right))
    for edge_i in ['top', 'bottom', 'right', 'left']:
        ax.spines[edge_i].set_edgecolor("white")

    thresh = matrix.max() / 2.
    for i, j in itertools.product(range(matrix.shape[0]), range(matrix.shape[1])):
        num = '{:.2f}'.format(matrix[i, j]) if normalize else int(matrix[i, j])
        plt.text(j, i, num,
                 verticalalignment='center',
                 horizontalalignment="center",
                 color="white" if num > thresh else "black")

    plt.ylabel('Prediction')
    plt.xlabel('True Label')

    plt.tight_layout()
    # plt.savefig('indicator/method_2.png', transparent=False, dpi=800)

    plt.show()


if __name__ == "__main__":
    trans_mat = np.array([[62, 16, 32, 9, 36],
                          [16, 16, 13, 8, 7],
                          [28, 16, 61, 8, 18],
                          [16, 2, 10, 40, 48],
                          [52, 11, 49, 8, 39]], dtype=int)
    label = ["Patt {}".format(i) for i in range(trans_mat.shape[0])]
    confused_matrix = np.asarray([[4, 4, 2, 1], [4, 3, 2, 0]])
    plot_confusion_matrix(classes=label,confused_matrix=confused_matrix)
